#include <unittest/unittest.h>

#include <cusp/array2d.h>
#include <cusp/multiply.h>
#include <cusp/qr.h>


template<typename MatrixType, typename MemorySpace>
void QRtestHelper(MatrixType & Ah, MemorySpace)
{
	cusp::array2d<float, MemorySpace> Q, R, A = Ah;
	cusp::qr(A, Q, R);

	cusp::array2d<float, cusp::host_memory> Qh = Q, Rh = R;
	
	/* check that R is upper triangular */
	for (size_t i = 0; i < Rh.num_rows; ++i)
		for (size_t j = 0; j < i; ++j)
			ASSERT_EQUAL(Rh(i,j), 0.0);
	
	
	/* check decomposition */
	cusp::array2d<float, cusp::host_memory> Ch;
	cusp::multiply(Qh, Rh, Ch);
	ASSERT_EQUAL(Ch.num_rows, Ah.num_rows);
	ASSERT_EQUAL(Ch.num_cols, Ah.num_cols);
	for (size_t i = 0; i < Ch.num_rows; ++i)
		for (size_t j = 0; j < Ch.num_cols; ++j)
			ASSERT_ALMOST_EQUAL(Ah(i,j), Ch(i,j));
}

template <class MemorySpace>
void TestQR(void)
{
	/* test with more rows than columns */

	cusp::array2d<float, cusp::host_memory> Ah(5,3);
	Ah(0,0) = 12.0; Ah(0,1) = -51.0;  Ah(0,2) =   4.0;
	Ah(1,0) =  6.0; Ah(1,1) = 167.0;  Ah(1,2) = -68.0;
	Ah(2,0) = -4.0; Ah(2,1) =  24.0;  Ah(2,2) = -41.0;
	Ah(3,0) =  8.0; Ah(3,1) =  45.0;  Ah(3,2) =   0.0;
	Ah(4,0) = -4.0; Ah(4,1) = -86.0;  Ah(4,2) =   1.0;

	QRtestHelper(Ah, MemorySpace());
}
DECLARE_HOST_DEVICE_UNITTEST(TestQR);

template <class MemorySpace>
void TestQR2(void)
{
	/* test with more columns than rows */

	cusp::array2d<float, cusp::host_memory> Ah(3,5);
	Ah(0,0) = 12.0; Ah(1,0) = -51.0;  Ah(2,0) =   4.0;
	Ah(0,1) =  6.0; Ah(1,1) = 167.0;  Ah(2,1) = -68.0;
	Ah(0,2) = -4.0; Ah(1,2) =  24.0;  Ah(2,2) = -41.0;
	Ah(0,3) =  8.0; Ah(1,3) =  45.0;  Ah(2,3) =   0.0;
	Ah(0,4) = -4.0; Ah(1,4) = -86.0;  Ah(2,4) =   1.0;

	QRtestHelper(Ah, MemorySpace());
}
DECLARE_HOST_DEVICE_UNITTEST(TestQR2);

